package com.itx.gatewayx.filter;

import cn.hutool.core.util.StrUtil;
import com.itx.gatewayx.common.constants.SystemConstants;
import com.itx.gatewayx.config.IgnoreUrlsConfig;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Arrays;
import java.util.List;

import javax.annotation.PostConstruct;

/**
 * 自定义限流过滤器
 *
 * @author itx
 */
@Slf4j
@Component
@RequiredArgsConstructor
@ConditionalOnProperty(name = "spring.redis.enabled", havingValue = "true", matchIfMissing = true)
public class RateLimiterFilter implements GlobalFilter, Ordered {

    private final RedisTemplate<String, Object> redisTemplate;
    private final IgnoreUrlsConfig ignoreUrlsConfig;
    
    @Value("${gateway.rate-limit.default-limit:100}")
    private int defaultLimit;
    
    private List<String> ignoreUrls;

    @PostConstruct
    public void init() {
        this.ignoreUrls = ignoreUrlsConfig.getIgnoreUrls();
    }

    /**
     * 限流Lua脚本，使用令牌桶算法
     */
    private static final String RATE_LIMITER_SCRIPT = 
            "local key = KEYS[1]\n" +
            "local now = tonumber(ARGV[1])\n" +
            "local requested = tonumber(ARGV[2])\n" +
            "local limit = tonumber(ARGV[3])\n" +
            "local ttl = 60\n" + // 60秒窗口期
            "local lastRefill = tonumber(redis.call('hget', key, 'lastRefill') or '0')\n" +
            "local tokens = tonumber(redis.call('hget', key, 'tokens') or limit)\n" +
            "local timeSinceLast = now - lastRefill\n" +
            "if timeSinceLast > 0 then\n" +
            "    tokens = math.min(limit, tokens + timeSinceLast)\n" +
            "    redis.call('hset', key, 'lastRefill', now)\n" +
            "end\n" +
            "if tokens >= requested then\n" +
            "    redis.call('hset', key, 'tokens', tokens - requested)\n" +
            "    redis.call('expire', key, ttl)\n" +
            "    return 1\n" +
            "end\n" +
            "return 0";

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        String path = request.getURI().getPath();
        
        // 公共接口不进行限流
        if (shouldIgnore(path)) {
            return chain.filter(exchange);
        }
        
        // 获取用户ID或IP
        String identity = getUser(request);
        if (StrUtil.isBlank(identity)) {
            identity = getIpAddress(request);
        }
        
        // 构造限流键
        String rateLimitKey = SystemConstants.RATE_LIMITER_KEY + path + ":" + identity;
        
        // 执行限流
        boolean allowed = isAllowed(rateLimitKey, defaultLimit);
        
        if (allowed) {
            return chain.filter(exchange);
        } else {
            return tooManyRequests(exchange.getResponse());
        }
    }

    /**
     * 执行限流逻辑
     */
    private boolean isAllowed(String key, int limit) {
        try {
            Long currentTime = Instant.now().getEpochSecond(); // 当前时间（秒级）
            List<String> keys = List.of(key);
            Long result = redisTemplate.execute(
                    RedisScript.of(RATE_LIMITER_SCRIPT, Long.class),
                    keys,
                    currentTime.toString(), "1", String.valueOf(limit)
            );
            return result != null && result == 1L;
        } catch (Exception e) {
            log.error("限流异常: {}", e.getMessage());
            return true; // 出现异常时，不阻止请求
        }
    }
    
    /**
     * 从请求中获取用户ID
     */
    private String getUser(ServerHttpRequest request) {
        return request.getHeaders().getFirst("X-User-ID");
    }

    /**
     * 获取请求IP地址
     */
    private String getIpAddress(ServerHttpRequest request) {
        String ip = request.getHeaders().getFirst("X-Forwarded-For");
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddress().getHostString();
        }
        // 如果是多级代理，获取第一个IP
        if (ip != null && ip.contains(",")) {
            ip = ip.split(",")[0].trim();
        }
        return ip;
    }
    
    /**
     * 判断是否是忽略限流的URL
     */
    private boolean shouldIgnore(String path) {
        return ignoreUrls.stream().anyMatch(url -> path.matches(url.replace("/**", ".*")));
    }

    /**
     * 返回请求过多响应
     */
    private Mono<Void> tooManyRequests(ServerHttpResponse response) {
        response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
        String body = "{\"code\":429,\"message\":\"请求过于频繁，请稍后再试\"}";
        DataBuffer buffer = response.bufferFactory().wrap(body.getBytes(StandardCharsets.UTF_8));
        return response.writeWith(Mono.just(buffer));
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 2000;
    }
} 